package cn.tworice.auth.util;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

public class KeySaving {
    public static final String PUBLIC_KEY_FILE = "public.key";
    public static final String PRIVATE_KEY_FILE = "private.key";
    public static final String TOKEN_FILE = "token";
    public static final String ENCRYPTED_KEY_FILE = "encrypted.key";

    public static void main(String[] args) throws Exception {
        // 生成对称密钥
        SecretKey secretKey = generateSymmetricKey();

        // 生成RSA密钥对
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048); // 使用2048位密钥长度
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();

        // 加密对称密钥
        byte[] encryptedKey = encryptKey(secretKey, publicKey);

        // 将加密后的对称密钥保存到文件中
        saveKeyToFile(encryptedKey, ENCRYPTED_KEY_FILE);

        // 从文件中读取加密后的对称密钥
        byte[] loadedKey = loadKeyFromFile(ENCRYPTED_KEY_FILE);

        // 解密对称密钥
        SecretKey decryptedKey = decryptKey(loadedKey, privateKey);

        // 输出解密后的对称密钥
        System.out.println("Decrypted Key: " + decryptedKey);
    }

    public static SecretKey generateSymmetricKey() throws Exception {
        KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
        keyGenerator.init(256); // 使用256位密钥长度
        return keyGenerator.generateKey();
    }

    public static byte[] encryptKey(SecretKey secretKey, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return cipher.doFinal(secretKey.getEncoded());
    }

    public static SecretKey decryptKey(byte[] encryptedKey, PrivateKey privateKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] decryptedKeyBytes = cipher.doFinal(encryptedKey);
        return new SecretKeySpec(decryptedKeyBytes, "AES");
    }

    public static void saveKeyToFile(byte[] keyData, String fileName) throws Exception {
        FileOutputStream fos = new FileOutputStream(fileName);
        fos.write(keyData);
        fos.close();
    }

    public static byte[] loadKeyFromFile(String fileName) throws Exception {
        return Files.readAllBytes(Paths.get(fileName));
    }
}
